#include "tcp_com.h"

#include <stdio.h>
#include <string.h>

#include "cfg.h"
#include "cmsis_os.h"
#include "ip.h"
#include "lwip.h"
#include "lwip/api.h"
#include "lwip/inet.h"
#include "lwip/init.h"
#include "lwip/ip.h"
#include "lwip/ip_addr.h"
#include "lwip/netif.h"
#include "lwip/opt.h"
#include "lwip/sockets.h"
#include "lwip/tcp.h"

typedef struct {
    uint8_t buffer[BUFFER_SIZE];
    uint16_t length;
    uint32_t last_rx_time;
    uint8_t timeout_flag;
} rx_buffer_t;

static rx_buffer_t rx_buffer = {0};

uint16_t check_frame(uint8_t *buffer, uint16_t length) {
    if (length < 9) {
        return 0;
    }

    if (buffer[0] != 0xAA || buffer[1] != 0x55 || buffer[2] != 0xAA ||
        buffer[3] != 0x55) {
        return 0;
    }

    uint16_t data_length = (buffer[7] << 8) | buffer[8];

    uint16_t total_length = 9 + data_length + 2;

    if (length < total_length) {
        return 0;
    }

    return total_length;
}

void process_tcp_data(uint8_t *new_data, uint16_t new_length, dev_send_ack cb,
                      int index) {
    uint32_t current_time = xTaskGetTickCount() * portTICK_PERIOD_MS;

    if (rx_buffer.length > 0 &&
        (current_time - rx_buffer.last_rx_time) > CMD_TIMEOUT_MS) {
        rx_buffer.length = 0;
        rx_buffer.timeout_flag = 1;
        printf("Timeout detected, buffer cleared\r\n");
    } else {
        rx_buffer.timeout_flag = 0;
    }

    rx_buffer.last_rx_time = current_time;

#if 0  // device busy ctrl
		static uint32_t last_response_time = 0;
    if (current_time - last_response_time > CMD_TIMEOUT_MS) {
        last_response_time = current_time;
    } else {
				// device busy
				uint8_t buff[12];
				buff[0] = new_data[0];
				buff[1] = new_data[1];
				buff[2] = new_data[2];
				buff[3] = new_data[3];
				buff[4] = new_data[4];
				buff[5] = new_data[5];
				buff[6] = new_data[6] | 0x80;
				buff[7] = 0x0;
				buff[8] = 0x1;
				buff[9] = DEVICE_BUSY;
				uint16_t crc_v = CRC16_Calculate(&buff[4], 6);
				buff[10] = crc_v & 0xFF;
				buff[11] = crc_v >> 8;
				cb(buff, 12);
				return;
    }
#endif
    if (rx_buffer.length + new_length <= BUFFER_SIZE) {
        memcpy(rx_buffer.buffer + rx_buffer.length, new_data, new_length);
        rx_buffer.length += new_length;
    } else {
        rx_buffer.length = 0;
        printf("Buffer overflow, buffer cleared\n");
        return;
    }

    uint16_t frame_length;
    uint16_t processed = 0;

    while ((frame_length = check_frame(rx_buffer.buffer + processed,
                                       rx_buffer.length - processed)) > 0) {
        cmd_process_detail(rx_buffer.buffer + processed, frame_length, cb,
                           index);
        processed += frame_length;
    }

    if (processed > 0) {
        if (processed < rx_buffer.length) {
            memmove(rx_buffer.buffer, rx_buffer.buffer + processed,
                    rx_buffer.length - processed);
            rx_buffer.length -= processed;
        } else {
            rx_buffer.length = 0;
        }
    }
}

struct netconn *srv_conn = NULL;
// struct netconn *srv_cli_conn = NULL;
struct netconn *cli_conn = NULL;

int need_reinit_server = 0;
// Define maximum number of clients
#define MAX_CLIENTS 2

// Client connection array and status flags
struct netconn *srv_cli_conn[MAX_CLIENTS] = {NULL, NULL};
int client_connected[MAX_CLIENTS] = {0, 0};

void ethernetif_notify_conn_changed(struct netif *netif) {
    if (!netif_is_link_up(netif)) {
        need_reinit_server = 1;
        for (int i = 0; i < MAX_CLIENTS; i++) {
            if (srv_cli_conn[i] != NULL) {
                printf("Closing server cli socket %d\r\n", i);
                netconn_close(srv_cli_conn[i]);
                netconn_delete(srv_cli_conn[i]);
                srv_cli_conn[i] = NULL;
            }
            client_connected[i] = 0;
        }

        if (srv_conn != NULL) {
            printf("Closing server socket\r\n");
            netconn_close(srv_conn);
            netconn_delete(srv_conn);
            srv_conn = NULL;
        }
    }
}

// Modified acknowledgment function to handle multiple clients
void tcp_srv_send_ack(uint8_t *data, uint16_t len, int client_index) {
    err_t err;

    if (client_index < 0 || client_index >= MAX_CLIENTS ||
        srv_cli_conn[client_index] == NULL) {
        printf("[srv] Invalid client index for send: %d\r\n", client_index);
        return;
    }

    err = netconn_write(srv_cli_conn[client_index], data, len, NETCONN_COPY);
    if (err != ERR_OK) {
        printf("[srv] Failed to send to client %d: %d\r\n", client_index, err);
    }
}

void StartTCPServerTask(void const *argument) {
    struct netbuf *buf;
    char *data;
    u16_t len;
    err_t err;

    while (1) {
        if (srv_conn == NULL || need_reinit_server) {
            if (srv_conn != NULL) {
                netconn_close(srv_conn);
                netconn_delete(srv_conn);
                srv_conn = NULL;
            }

            // Close and clean up any existing client connections
            for (int i = 0; i < MAX_CLIENTS; i++) {
                if (srv_cli_conn[i] != NULL) {
                    netconn_close(srv_cli_conn[i]);
                    netconn_delete(srv_cli_conn[i]);
                    srv_cli_conn[i] = NULL;
                }
                client_connected[i] = 0;
            }

            vTaskDelay(1000);

            srv_conn = netconn_new(NETCONN_TCP);
            if (srv_conn == NULL) {
                printf("[srv] Failed to create server netconn\r\n");
                vTaskDelay(1000);
                continue;
            }

            netconn_set_nonblocking(srv_conn, 0);
            ip_set_option(srv_conn->pcb.ip, SOF_REUSEADDR);

            err = netconn_bind(srv_conn, IP_ADDR_ANY, cfg.port);
            if (err != ERR_OK) {
                printf("[srv] Bind failed: %d\r\n", err);
                netconn_delete(srv_conn);
                srv_conn = NULL;
                vTaskDelay(1000);
                continue;
            }

            err = netconn_listen(srv_conn);
            if (err != ERR_OK) {
                printf("[srv] Listen failed: %d\r\n", err);
                netconn_delete(srv_conn);
                srv_conn = NULL;
                vTaskDelay(1000);
                continue;
            }

            printf("[srv] TCP server listening on port %d\r\n", cfg.port);
            need_reinit_server = 0;
        }

        netconn_set_recvtimeout(srv_conn, 10);

        // Check for new client connections
        struct netconn *new_client = NULL;
        while (srv_conn && need_reinit_server == 0) {
            err = netconn_accept(srv_conn, &new_client);
            if (err == ERR_OK && new_client != NULL) {
                // Find an available slot for the new client
                int slot = -1;
                for (int i = 0; i < MAX_CLIENTS; i++) {
                    if (srv_cli_conn[i] == NULL) {
                        slot = i;
                        break;
                    }
                }

                if (slot >= 0) {
                    // We have a free slot
                    srv_cli_conn[slot] = new_client;
                    client_connected[slot] = 1;

                    printf("[srv] Client %d connected\r\n", slot);

                    netconn_set_recvtimeout(srv_cli_conn[slot], 10);

                    struct tcp_pcb *pcb = srv_cli_conn[slot]->pcb.tcp;
                    if (pcb) {
                        ip_set_option(pcb, SOF_KEEPALIVE);
                        pcb->keep_idle = 3000;
                        pcb->keep_intvl = 500;
                        pcb->keep_cnt = 3;
                    }
                } else {
                    // No slots available, reject connection
                    printf(
                        "[srv] Max clients reached, rejecting connection\r\n");
                    netconn_close(new_client);
                    netconn_delete(new_client);
                }
            } else if (err != ERR_TIMEOUT && err != ERR_OK) {
                printf("[srv] Accept failed: %d, reinitializing server\r\n",
                       err);
                need_reinit_server = 1;
                continue;
            }

            // Process data from connected clients
            for (int i = 0; i < MAX_CLIENTS; i++) {
                if (srv_cli_conn[i] != NULL && client_connected[i]) {
                    if (need_reinit_server) {
                        break;
                    }

                    // Pass client index to send function for proper client
                    // identification
                    err = netconn_recv(srv_cli_conn[i], &buf);
                    if (err == ERR_OK) {
                        printf("[srv] recv client %d cmd\r\n", i);
                        netbuf_data(buf, (void **)&data, &len);

                        // Modified to pass client index to the processing
                        // function
                        process_tcp_data((uint8_t *)data, len, tcp_srv_send_ack,
                                         i);

                        netbuf_delete(buf);
                    } else if (err == ERR_TIMEOUT) {
                        continue;
                    } else {
                        printf("[srv] Client %d disconnected, error: %d\r\n", i,
                               err);
                        netconn_close(srv_cli_conn[i]);
                        netconn_delete(srv_cli_conn[i]);
                        srv_cli_conn[i] = NULL;
                        client_connected[i] = 0;
                    }
                }
            }

            vTaskDelay(10);
        }
    }
}

void tcp_cli_send_ack(uint8_t *data, uint16_t len, int index) {
    if (cli_conn == NULL) {
        printf("[cli] client cli_conn err\r\n");
        return;
    }
    struct netbuf *buf = netbuf_new();
    if (buf == NULL) {
        printf("[cli] netbuf_new failed\r\n");
        return;
    }
    void *payload = netbuf_alloc(buf, len);
    if (payload == NULL) {
        printf("[cli] netbuf_alloc failed\r\n");
        netbuf_delete(buf);
        return;
    }
    memcpy(payload, data, len);
    netconn_write(cli_conn, payload, len, NETCONN_COPY);
    printf("[cli] send ack to server\r\n");
    netbuf_delete(buf);
}

void StartTCPClientTask(void const *argument) {
    printf("StartTCPClientTask\r\n");

    struct netbuf *buf;
    ip_addr_t server_ip;
    err_t err;
    int link_down_detected = 0;

    while (1) {
        if (netif_default == NULL || !netif_is_up(netif_default) ||
            !netif_is_link_up(netif_default)) {
            if (cli_conn != NULL) {
                printf("[cli] Network link down, closing connection\r\n");
                netconn_close(cli_conn);
                netconn_delete(cli_conn);
                cli_conn = NULL;
            }
            link_down_detected = 1;
            osDelay(1000);
            continue;
        }

        if (link_down_detected && netif_is_link_up(netif_default)) {
            printf("[cli] Network link restored\r\n");
            link_down_detected = 0;
        }

        if (cli_conn == NULL) {
            cli_conn = netconn_new(NETCONN_TCP);
            if (cli_conn == NULL) {
                printf("[cli] Client: Create netconn failed!\r\n");
                osDelay(1000);
                continue;
            }

            netconn_set_recvtimeout(cli_conn, 2000);

            IP4_ADDR(&server_ip, cfg.dest_ip[0], cfg.dest_ip[1], cfg.dest_ip[2],
                     cfg.dest_ip[3]);
            printf("[cli] Connecting to server %d.%d.%d.%d:%d\r\n",
                   cfg.dest_ip[0], cfg.dest_ip[1], cfg.dest_ip[2],
                   cfg.dest_ip[3], cfg.dest_port);

            err = netconn_connect(cli_conn, &server_ip, cfg.dest_port);
            if (err != ERR_OK) {
                printf("[cli] Connect failed: %d\r\n", err);
                netconn_delete(cli_conn);
                cli_conn = NULL;
                osDelay(3000);
                continue;
            }

            struct tcp_pcb *pcb = cli_conn->pcb.tcp;
            if (pcb) {
                ip_set_option(pcb, SOF_KEEPALIVE);
                pcb->keep_idle = 3000;
                pcb->keep_intvl = 500;
                pcb->keep_cnt = 3;
            }

            printf("[cli] Connected to server\r\n");
        }

        err = netconn_recv(cli_conn, &buf);
        if (err == ERR_OK) {
            uint16_t len;
            void *data;
            netbuf_data(buf, &data, &len);
            process_tcp_data((uint8_t *)data, len, tcp_cli_send_ack, 0);
            netbuf_delete(buf);
        } else if (err == ERR_TIMEOUT) {
            if (!netif_is_link_up(netif_default)) {
                printf("[cli] Link down detected during timeout\r\n");
                if (cli_conn != NULL) {
                    netconn_close(cli_conn);
                    netconn_delete(cli_conn);
                    cli_conn = NULL;
                }
                link_down_detected = 1;
            }
            osDelay(100);
        } else {
            printf("[cli] Receive failed: %d\r\n", err);
            if (cli_conn != NULL) {
                netconn_close(cli_conn);
                netconn_delete(cli_conn);
                cli_conn = NULL;
            }
            osDelay(3000);
        }
    }
}
